package com.bruce.ai.alibaba.agentcard.discovery.client;

import com.alibaba.nacos.api.ai.model.a2a.AgentCard;
import com.bruce.ai.alibaba.agentcard.discovery.circuitbreaker.CircuitBreaker;
import com.bruce.ai.alibaba.agentcard.discovery.circuitbreaker.CircuitBreakerConfig;
import com.bruce.ai.alibaba.agentcard.discovery.retry.RetryExecutor;
import com.bruce.ai.alibaba.agentcard.discovery.retry.RetryStrategy;
import com.bruce.ai.alibaba.agentcard.discovery.config.A2AClientConfig;
import com.bruce.ai.alibaba.agentcard.discovery.model.A2ARequest;
import com.bruce.ai.alibaba.agentcard.discovery.model.A2AResponse;
import com.bruce.ai.alibaba.agentcard.discovery.metrics.A2AMetrics;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.time.Instant;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * A2A客户端实现，用于与Agent进行通信
 * 基于Google A2A协议标准和Nacos 3.1的A2A功能实现
 * 
 * 特性：
 * - HTTP/HTTPS通信
 * - 重试机制（指数退避）
 * - 熔断器模式
 * - 异步调用支持
 * - 多种认证方式
 * - 连接池管理
 */
public class A2AClient {
    
    private static final Logger logger = LoggerFactory.getLogger(A2AClient.class);
    private static final MediaType JSON = MediaType.get("application/json; charset=utf-8");
    
    private final AgentCard agentCard;
    private final A2AClientConfig config;
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private final CircuitBreaker circuitBreaker;
    private final String sourceAgentName;
    private final A2AMetrics metrics;
    
    /**
     * 构造函数
     * @param agentCard Nacos AgentCard
     * @param config A2A客户端配置
     * @param sourceAgentName 源Agent名称
     */
    public A2AClient(AgentCard agentCard, A2AClientConfig config, String sourceAgentName) {
        this.agentCard = agentCard;
        this.config = config;
        this.sourceAgentName = sourceAgentName;
        this.httpClient = createHttpClient();
        this.objectMapper = createObjectMapper();
        this.metrics = new A2AMetrics();
        this.circuitBreaker = config.isCircuitBreakerEnabled() ? 
            new CircuitBreaker(
                agentCard.getName(), 
                CircuitBreakerConfig.builder()
                    .failureThreshold(config.getCircuitBreakerFailureThreshold())
                    .waitDurationInOpenState(java.time.Duration.ofMillis(config.getCircuitBreakerRecoveryTimeoutMs()))
                    .build()
            ) : null;
        
        logger.info("A2AClient initialized for agent: {} with endpoint: {}", 
                    agentCard.getName(), agentCard.getUrl());
    }
    
    /**
     * 同步调用Agent
     * @param message 要发送的消息
     * @return Agent的响应
     */
    public A2AResponse call(String message) {
        return call(new A2ARequest(sourceAgentName, agentCard.getName(), message));
    }
    
    /**
     * 同步调用Agent
     * @param request A2A请求对象
     * @return Agent的响应
     */
    public A2AResponse call(A2ARequest request) {
        try {
            logger.info("Calling agent {} with request: {}", agentCard.getName(), request.getId());
            
            if (circuitBreaker != null) {
                return circuitBreaker.execute(() -> executeRequest(request));
            } else {
                return executeRequest(request);
            }
        } catch (Exception e) {
            logger.error("Error calling agent {}", agentCard.getName(), e);
            throw new RuntimeException("Failed to call agent: " + e.getMessage(), e);
        }
    }
    
    /**
     * 异步调用Agent
     * @param message 要发送的消息
     * @return 异步响应Future
     */
    public CompletableFuture<A2AResponse> callAsync(String message) {
        return callAsync(new A2ARequest(sourceAgentName, agentCard.getName(), message));
    }
    
    /**
     * 异步调用Agent
     * @param request A2A请求对象
     * @return 异步响应Future
     */
    public CompletableFuture<A2AResponse> callAsync(A2ARequest request) {
        return CompletableFuture.supplyAsync(() -> {
            try {
                logger.info("Async calling agent {} with request: {}", agentCard.getName(), request.getId());
                return executeRequest(request);
            } catch (Exception e) {
                logger.error("Error in async call to agent {}", agentCard.getName(), e);
                throw new RuntimeException("Failed to call agent asynchronously: " + e.getMessage(), e);
            }
        }).exceptionally(throwable -> {
            // 异步异常处理
            logger.error("Async call to agent {} failed with exception", agentCard.getName(), throwable);
            
            A2AResponse errorResponse = new A2AResponse();
            errorResponse.setRequestId(request.getId());
            errorResponse.setSourceAgent(agentCard.getName());
            errorResponse.setTargetAgent(request.getSourceAgent());
            errorResponse.setStatus("ERROR");
            errorResponse.setErrorCode("ASYNC_EXECUTION_FAILED");
            errorResponse.setErrorMessage("Async execution failed: " + throwable.getMessage());
            errorResponse.setTimestamp(Instant.now());
            
            return errorResponse;
        });
    }
    
    /**
     * 异步调用Agent（带超时）
     * @param request A2A请求对象
     * @param timeoutMs 超时时间（毫秒）
     * @return 异步响应Future
     */
    public CompletableFuture<A2AResponse> callAsync(A2ARequest request, long timeoutMs) {
        CompletableFuture<A2AResponse> future = callAsync(request);
        
        return future.orTimeout(timeoutMs, TimeUnit.MILLISECONDS)
                .exceptionally(throwable -> {
                    if (throwable instanceof java.util.concurrent.TimeoutException) {
                        logger.warn("Async call to agent {} timed out after {}ms", agentCard.getName(), timeoutMs);
                        
                        A2AResponse timeoutResponse = new A2AResponse();
                        timeoutResponse.setRequestId(request.getId());
                        timeoutResponse.setSourceAgent(agentCard.getName());
                        timeoutResponse.setTargetAgent(request.getSourceAgent());
                        timeoutResponse.setStatus("ERROR");
                        timeoutResponse.setErrorCode("TIMEOUT");
                        timeoutResponse.setErrorMessage("Request timed out after " + timeoutMs + "ms");
                        timeoutResponse.setTimestamp(Instant.now());
                        
                        return timeoutResponse;
                    } else {
                        // 重新抛出其他异常
                        throw new RuntimeException(throwable);
                    }
                });
    }
    
    /**
     * 获取Agent卡片信息
     * @return AgentCard
     */
    public AgentCard getAgentCard() {
        return agentCard;
    }
    
    /**
     * 获取指标收集器
     * @return A2AMetrics
     */
    public A2AMetrics getMetrics() {
        return metrics;
    }
    
    /**
     * 获取指标摘要
     * @return MetricsSummary
     */
    public A2AMetrics.MetricsSummary getMetricsSummary() {
        return metrics.getSummary();
    }
    
    /**
     * 检查Agent是否可用
     * @return 是否可用
     */
    public boolean isAvailable() {
        try {
            // 发送健康检查请求
            A2ARequest healthCheck = new A2ARequest(sourceAgentName, agentCard.getName(), "health_check");
            healthCheck.setMessageType("HEALTH_CHECK");
            healthCheck.setTimeoutMs(config.getHealthCheckTimeoutMs());
            
            A2AResponse response = executeRequest(healthCheck);
            return "SUCCESS".equals(response.getStatus());
        } catch (Exception e) {
            logger.warn("Agent {} is not available: {}", agentCard.getName(), e.getMessage());
            return false;
        }
    }
    
    /**
     * 执行HTTP请求
     * @param request A2A请求
     * @return A2A响应
     */
    private A2AResponse executeRequest(A2ARequest request) throws IOException {
        long startTime = System.currentTimeMillis();
        String agentName = agentCard.getName();
        
        // 记录请求开始
        metrics.recordRequestStart(agentName);
        
        try {
            // 如果启用了熔断器，使用熔断器执行请求
            if (circuitBreaker != null) {
                try {
                    A2AResponse response = circuitBreaker.execute(() -> {
                        return executeRequestInternal(request, startTime);
                    }, (exception) -> {
                        // 熔断器开启时的回退响应
                        metrics.recordCircuitBreakerOpen(agentName);
                        
                        A2AResponse fallbackResponse = new A2AResponse();
                        fallbackResponse.setRequestId(request.getId());
                        fallbackResponse.setSourceAgent(agentCard.getName());
                        fallbackResponse.setTargetAgent(request.getSourceAgent());
                        fallbackResponse.setStatus("ERROR");
                        fallbackResponse.setErrorCode("CIRCUIT_BREAKER_OPEN");
                        fallbackResponse.setErrorMessage("Circuit breaker is open: " + exception.getMessage());
                        fallbackResponse.setTimestamp(Instant.now());
                        fallbackResponse.setProcessingTimeMs(System.currentTimeMillis() - startTime);
                        return fallbackResponse;
                    });
                    
                    // 设置处理时间
                    long responseTime = System.currentTimeMillis() - startTime;
                    if (response.getProcessingTimeMs() == null) {
                        response.setProcessingTimeMs(responseTime);
                    }
                    
                    return response;
                    
                } catch (Exception e) {
                    long responseTime = System.currentTimeMillis() - startTime;
                    metrics.recordRequestFailure(agentName, "CIRCUIT_BREAKER_EXECUTION_FAILED", responseTime, 0);
                    throw new IOException("Circuit breaker execution failed", e);
                }
            } else {
                // 没有熔断器，直接执行
                A2AResponse response = executeRequestInternal(request, startTime);
                long responseTime = System.currentTimeMillis() - startTime;
                response.setProcessingTimeMs(responseTime);
                
                return response;
            }
        } catch (IOException e) {
            long responseTime = System.currentTimeMillis() - startTime;
            metrics.recordRequestFailure(agentName, "IO_EXCEPTION", responseTime, 0);
            throw e;
        } catch (Exception e) {
            long responseTime = System.currentTimeMillis() - startTime;
            metrics.recordRequestFailure(agentName, "UNKNOWN_EXCEPTION", responseTime, 0);
            throw new IOException("Unexpected error during request execution", e);
        }
    }
    
    /**
     * 内部请求执行方法
     */
    private A2AResponse executeRequestInternal(A2ARequest request, long startTime) throws IOException {
        // 构建请求URL
        String url = buildRequestUrl();
        
        // 序列化请求体
        String requestBody = objectMapper.writeValueAsString(request);
        
        // 构建HTTP请求
        Request.Builder requestBuilder = new Request.Builder()
                .url(url)
                .post(RequestBody.create(requestBody, JSON))
                .addHeader("Content-Type", "application/json")
                .addHeader("User-Agent", "A2AClient/1.0");
        
        // 添加认证头
        addAuthenticationHeaders(requestBuilder);
        
        Request httpRequest = requestBuilder.build();
        
        // 执行请求（带重试）
        return executeWithRetry(httpRequest, request, startTime);
    }
    
    /**
     * 带重试的请求执行
     */
    private A2AResponse executeWithRetry(Request httpRequest, A2ARequest originalRequest, long startTime) throws IOException {
        RetryStrategy retryStrategy = new RetryStrategy.Builder()
                .maxRetries(config.getMaxRetries())
                .baseDelay(config.getRetryBaseDelayMs())
                .maxDelay(config.getRetryMaxDelayMs())
                .algorithm(RetryStrategy.RetryAlgorithm.EXPONENTIAL_BACKOFF)
                .build();
        
        RetryExecutor retryExecutor = new RetryExecutor(retryStrategy);
        
        // 用于跟踪重试次数
        final int[] retryCount = {0};
        
        try {
            A2AResponse response = retryExecutor.execute(() -> {
                try (Response httpResponse = httpClient.newCall(httpRequest).execute()) {
                    if (httpResponse.isSuccessful()) {
                        String responseBody = httpResponse.body().string();
                        A2AResponse a2aResponse = objectMapper.readValue(responseBody, A2AResponse.class);
                        
                        logger.info("Successfully called agent {}", agentCard.getName());
                        return a2aResponse;
                    } else {
                        throw new IOException("HTTP " + httpResponse.code() + ": " + httpResponse.message());
                    }
                }
            }, (attempt, exception, delayMs) -> {
                retryCount[0]++; // 每次重试都增加计数
                logger.warn("Attempt {} failed for agent {}: {}, retrying in {}ms", 
                           attempt, agentCard.getName(), exception.getMessage(), delayMs);
            });
            
            // 计算响应时间并记录成功的指标
            long responseTime = System.currentTimeMillis() - startTime;
            metrics.recordRequestSuccess(agentCard.getName(), responseTime, retryCount[0]);
            return response;
            
        } catch (Exception e) {
            // 计算响应时间并记录失败的指标
            long responseTime = System.currentTimeMillis() - startTime;
            metrics.recordRequestFailure(agentCard.getName(), "RETRY_EXHAUSTED", responseTime, retryCount[0]);
            
            // 所有重试都失败了，返回错误响应
            A2AResponse errorResponse = new A2AResponse();
            errorResponse.setRequestId(originalRequest.getId());
            errorResponse.setSourceAgent(agentCard.getName());
            errorResponse.setTargetAgent(originalRequest.getSourceAgent());
            errorResponse.setStatus("ERROR");
            errorResponse.setErrorCode("RETRY_EXHAUSTED");
            errorResponse.setErrorMessage("All retry attempts failed: " + e.getMessage());
            errorResponse.setTimestamp(Instant.now());
            errorResponse.setProcessingTimeMs(responseTime);
            
            return errorResponse;
        }
    }
    

    
    /**
     * 构建请求URL
     */
    private String buildRequestUrl() {
        String endpoint = agentCard.getUrl();
        if (!endpoint.startsWith("http://") && !endpoint.startsWith("https://")) {
            endpoint = "http://" + endpoint;
        }
        
        if (!endpoint.endsWith("/")) {
            endpoint += "/";
        }
        
        return endpoint + "a2a/call";
    }
    
    /**
     * 添加认证头
     */
    private void addAuthenticationHeaders(Request.Builder requestBuilder) {
        switch (config.getAuthType()) {
            case BEARER:
                if (StringUtils.hasText(config.getBearerToken())) {
                    requestBuilder.addHeader("Authorization", "Bearer " + config.getBearerToken());
                }
                break;
            case API_KEY:
                if (StringUtils.hasText(config.getApiKey())) {
                    requestBuilder.addHeader("X-API-Key", config.getApiKey());
                }
                break;
            case OAUTH2:
                if (StringUtils.hasText(config.getOauth2Token())) {
                    requestBuilder.addHeader("Authorization", "Bearer " + config.getOauth2Token());
                }
                break;
            case NONE:
            default:
                // 无认证
                break;
        }
    }
    
    /**
     * 创建HTTP客户端
     */
    private OkHttpClient createHttpClient() {
        ConnectionPool connectionPool = new ConnectionPool(
                config.getMaxIdleConnections(),
                config.getKeepAliveDurationMs(),
                TimeUnit.MILLISECONDS
        );
        
        return new OkHttpClient.Builder()
                .connectTimeout(config.getConnectTimeoutMs(), TimeUnit.MILLISECONDS)
                .readTimeout(config.getReadTimeoutMs(), TimeUnit.MILLISECONDS)
                .writeTimeout(config.getWriteTimeoutMs(), TimeUnit.MILLISECONDS)
                .connectionPool(connectionPool)
                .retryOnConnectionFailure(false) // 我们自己处理重试
                .build();
    }
    
    /**
     * 创建JSON对象映射器
     */
    private ObjectMapper createObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.registerModule(new JavaTimeModule());
        return mapper;
    }
    
    /**
     * 关闭客户端资源
     */
    public void close() {
        if (httpClient != null) {
            httpClient.dispatcher().executorService().shutdown();
            httpClient.connectionPool().evictAll();
        }
    }
}